package top.zhangjianyong.tools.indicator.strategy.impl;

import org.springframework.stereotype.Component;
import top.zhangjianyong.tools.entity.ETFData;
import top.zhangjianyong.tools.indicator.context.IndicatorContext;
import top.zhangjianyong.tools.indicator.model.TechnicalIndicatorResult;

import java.math.BigDecimal;
import java.math.RoundingMode;
import java.util.List;

/**
 * 动量指标计算器
 * 计算5日、10日、20日收益率（mom_5、mom_10、mom_20）
 *
 * @author zhangjianyong
 * @date 2024/01/01
 */
@Component
public class MomentumIndicator extends AbstractIndicator {

    private static final String INDICATOR_NAME = "MOMENTUM";

    @Override
    public String getName() {
        return INDICATOR_NAME;
    }

    @Override
    public int getMinDataSize() {
        return 20; // 最大周期是20日
    }

    @Override
    protected TechnicalIndicatorResult doCalculate(IndicatorContext context) {
        List<ETFData> dataList = context.getDataList();

        if (dataList.isEmpty()) {
            return TechnicalIndicatorResult.failure(getName(), "数据列表为空");
        }

        BigDecimal currentPrice = dataList.get(dataList.size() - 1).getValue();

        TechnicalIndicatorResult result = TechnicalIndicatorResult.success(getName());

        // 计算5日收益率
        BigDecimal mom5 = calculateMomentum(dataList, currentPrice, 5);
        if (mom5 != null) {
            result.addValue("mom_5", mom5.setScale(4, RoundingMode.HALF_UP));
        }

        // 计算10日收益率
        BigDecimal mom10 = calculateMomentum(dataList, currentPrice, 10);
        if (mom10 != null) {
            result.addValue("mom_10", mom10.setScale(4, RoundingMode.HALF_UP));
        }

        // 计算20日收益率
        BigDecimal mom20 = calculateMomentum(dataList, currentPrice, 20);
        if (mom20 != null) {
            result.addValue("mom_20", mom20.setScale(4, RoundingMode.HALF_UP));
        }

        return result;
    }

    /**
     * 计算指定周期的收益率
     * 收益率 = (当前价格 - N日前价格) / N日前价格 * 100
     *
     * @param dataList    K线数据列表
     * @param currentPrice 当前价格
     * @param period      周期
     * @return 收益率（百分比）
     */
    private BigDecimal calculateMomentum(List<ETFData> dataList, BigDecimal currentPrice, int period) {
        if (dataList.size() < period + 1) {
            return null;
        }

        // 获取N日前的价格
        int previousIndex = dataList.size() - period - 1;
        BigDecimal previousPrice = dataList.get(previousIndex).getValue();

        if (previousPrice.compareTo(BigDecimal.ZERO) == 0) {
            return null;
        }

        // 计算收益率
        BigDecimal momentum = currentPrice.subtract(previousPrice)
            .divide(previousPrice, 4, RoundingMode.HALF_UP)
            .multiply(new BigDecimal(100));

        return momentum;
    }
}
